import MDAnalysis as mda
import numpy as np
import matplotlib.pyplot as plt
import os

def get_total_time(mdinfo_path):
    """Reads the total simulation time from the specified mdinfo file."""
    with open(mdinfo_path, 'r') as file:
        lines = file.readlines()
    # Extract time from the second line, which should contain 'TIME(PS) ='
    for line in lines:
        if 'TIME(PS) =' in line:
            # Split the line and extract the time value
            time_ps = float(line.split('TIME(PS) =')[1].split()[0])
            return time_ps
    raise ValueError(f"Could not find 'TIME(PS) =' in {mdinfo_path}")

# Choose the mdinfo file
mdinfo_path = r'v2_pro8_progress\rw_lf_pro8_trajout\pro9.mdinfo' if os.path.exists(r'v2_pro8_progress\rw_lf_pro8_trajout\pro9.mdinfo') else r'v2_pro8_progress\r375l_v2_pro9_trajout\pro8.mdinfo'

# Get the total simulation time
total_time_ps = get_total_time(mdinfo_path)
print(f'Total simulation time: {total_time_ps} ps')

# Load the trajectory and topology files
u = mda.Universe(r'v2_pro8_progress\rw_lf_pro8_trajout\traj.pdb', r'v2_pro8_progress\rw_lf_pro8_trajout\traj.dcd')

# Define the two regions (selection strings can be modified as needed)
region1 = u.select_atoms('resid 25:38')  # upper pore a
region2 = u.select_atoms('resid 326:335')  # upper pore C

# Initialize an array to store distances
distances = []

# Iterate through each frame of the trajectory
for ts in u.trajectory:
    com1 = region1.center_of_mass()
    com2 = region2.center_of_mass()
    distance = np.linalg.norm(com1 - com2)
    distances.append(distance)



# Convert distances to a numpy array
distances = np.array(distances)

# Calculate the overall average distance
average_distance = np.mean(distances)
print(f'Overall average distance: {average_distance:.3f} Å')

# Generate the time array based on the total simulation time, converting to ns
time = np.linspace(0, total_time_ps / 1000, len(distances))  # Convert from ps to ns

plt.xlim(0, 944)
plt.ylim(8, 18)

# Plot the distance vs. simulation time
plt.plot(time, distances, label='Distance', color='lightsteelblue')
plt.axhline(y=average_distance, color='r', linestyle='--', label='Average Distance')
plt.xlabel('Time (ns)', fontsize=16)
plt.ylabel('Distance (Å)', fontsize=16)
plt.xticks(fontsize=12)
plt.yticks(fontsize=12)
plt.title('RW_LF_a1a8_distance_chainA', fontsize=19, pad=20)
#plt.legend(loc='best')

# Save the plot to a file
plt.savefig(r'v2_pro8_progress\rw_lf_pro8_trajout\a1a8_distance_cA_RW_LF.svg', dpi=500)

plt.show()
]